import torch
from sklearn.metrics import classification_report, f1_score, accuracy_score, recall_score, precision_score

from __00__config import Config

config = Config()

def model2dev(model, data_loader, device):
	"""
	    在验证或测试集上评估 BERT 分类模型的性能。
    参数：
        model (nn.Module): BERT 分类模型。
        data_loader (DataLoader): 数据加载器（验证或测试集）。
        device (str): 设备（"cuda" 或 "cpu"）。
    返回：
        tuple: (分类报告, F1 分数, 准确度, 精确度，召回率)
            - report: 分类报告（包含每个类别的精确度、召回率、F1 分数等）。
            - f1score: 微平均 F1 分数。
            - accuracy: 准确度。
            - precision: 微平均精确度
            - recall: 微平均召回率
	"""
	# 将模型改为评估模式
	model.eval()
	# 初始化列表，存储预测结果和真实标签
	all_predicts, all_labels = [], []
	# 禁用梯度，提高效率
	with torch.no_grad():
		print("开始评估...")
		for i, batch in enumerate(data_loader):
			input_ids, attention_masks, labels = batch
			# 将数据移动到设备上
			input_ids = input_ids.to(device)
			attention_masks = attention_masks.to(device)
			labels = labels.to(device)
			# 前向传播
			output = model(input_ids, attention_masks)
			y_predict_list = torch.argmax(output, dim=-1)
			# 存储预测和真实标签
			all_predicts.extend(y_predict_list.cpu().tolist())
			all_labels.extend(labels.cpu().tolist())
	# 计算分类报告
	report = classification_report(all_labels, all_predicts)
	# f1分数
	f1 = f1_score(all_labels, all_predicts, average='macro')
	# 准确率
	accuracy = accuracy_score(all_labels, all_predicts)
	# 召回率
	recall = recall_score(all_labels, all_predicts, average='macro')
	# 精确率
	precision = precision_score(all_labels, all_predicts, average='macro')
	return report, f1, accuracy, precision, recall
